#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   crux_score.py
@Author :   Song
@Time   :   2022/5/19 16:54
@Contact:   songjian@westlake.edu.cn
@intro  :   给每个二级谱图添加Alpha-Filter的打分，然后作为后打分给crux
'''
import numpy as np
import pandas as pd
import click
import sys
from pathlib import Path
import models
import predifine
import torch
from make_npz import mzml_Reader
from dataloader import mzML_Dataset
from train import my_collate, test_one_epoch
import time

# @click.option('--dir_ws', required=False, help='Specify the target workspace')
def main():
    dir_ws = Path('/home/songjian/Alpha-Filter/huh7')
    dir_mzml = list(dir_ws.rglob('*.mzML'))[0]
    dir_input = dir_ws / 'crux-output' / 'make-pin.pin'
    dir_output = dir_ws / 'filter-output'
    dir_output.mkdir(exist_ok=True)
    dir_output = dir_output / 'make-pin.pin'
    dir_model = '/home/songjian/Alpha-Filter/train_output/Filter_epoch_1.pt'
    device = predifine.device

    # 读取模型
    model = models.Model_Filter(
        dim_model=predifine.dim_model,
        n_head=predifine.n_head,
        dim_feedforward=predifine.dim_feedforward,
        n_layers=predifine.n_layers,
        dropout=predifine.dropout,
        dim_intensity=predifine.dim_intensity,
        max_length=predifine.max_length,
        max_charge=predifine.max_charge
    )
    model.load_state_dict(torch.load(dir_model))
    model = model.to(device)

    # 读取mzml
    mzml = mzml_Reader(dir_mzml)

    # 读取crux
    df_crux = pd.read_csv(dir_input, sep='\t')
    test_idx = (df_crux['ScanNr'] - 1).unique()
    assert (mzml.all_levels[test_idx] == 2).all()

    # Alpha-Filter
    test_dataset = mzML_Dataset(mzml, test_idx)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=128,
                                              num_workers=predifine.num_workers,
                                              shuffle=False,
                                              pin_memory=True,
                                              collate_fn=my_collate)  # 默认collate_fn
    scores = test_one_epoch(test_loader, model, device)
    scores = pd.Series(scores)
    scores.index = test_idx
    scores = scores[df_crux['ScanNr'] - 1].values

    # save
    df_crux.insert(loc=10, column='filter_score', value=scores)
    df_crux.to_csv(dir_output, sep='\t', index=False)


if __name__ == '__main__':
    t0 = time.time()
    main()
    ut = (time.time() - t0) / 60.
    print(f'finished in {ut:.3f}')
